import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit


def build_trt_engine(onnx_path, engine_path, precision='fp16'):
    """
    构建 TensorRT 引擎
    :param precision: 'fp32', 'fp16', 'int8'
    """
    logger = trt.Logger(trt.Logger.WARNING)
    builder = trt.Builder(logger)
    network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
    parser = trt.OnnxParser(network, logger)

    # 加载 ONNX 模型
    with open(onnx_path, 'rb') as model:
        if not parser.parse(model.read()):
            print("❌ ONNX 解析错误:")
            for error in range(parser.num_errors):
                print(parser.get_error(error))
            return None

    # 配置构建选项
    config = builder.create_builder_config()
    # config.max_workspace_size = 4 * 1024 * 1024 * 1024  # 4GB

    # 精度设置
    if precision == 'fp16' and builder.platform_has_fast_fp16:
        config.set_flag(trt.BuilderFlag.FP16)
    elif precision == 'int8' and builder.platform_has_fast_int8:
        config.set_flag(trt.BuilderFlag.INT8)
        # 在此添加校准器 (见下文)

    # 构建引擎
    serialized_engine = builder.build_serialized_network(network, config)

    # 保存引擎
    with open(engine_path, 'wb') as f:
        f.write(serialized_engine)

    print(f"🚀 TensorRT {precision.upper()} 引擎构建成功: {engine_path}")
    return serialized_engine


# 示例使用
build_trt_engine('../models/yolov8n_optimized.onnx', 'yolov8n_fp16.engine', 'fp16')